#include <iostream>
#include <fstream>
#include "body.h"
#include "tinyobj.h"
#include "tinyply.h"

int BodyBuilder::body_count = 0;

std::shared_ptr<MeshData>
MeshData::LoadMeshFile(const std::filesystem::path& file_path) {
	auto ext = file_path.extension();
	std::ifstream stream(file_path, std::ios::binary);

	if (!stream.is_open()) {
		std::cerr << "[ERROR] fail to open mesh file: "
			<< file_path << std::endl;
		return nullptr;
	}

	auto mesh_ptr = std::make_shared<MeshData>();
	bool ret = false;
	if (ext == ".obj") {
		ret = mesh_ptr->LoadObjMesh(stream);
	}
	else if (ext == ".ply") {
		ret = mesh_ptr->LoadPlyMesh(stream);
	}

	if (!ret) {
		std::cerr << "[ERROR] fail to parse mesh file: "
			<< file_path << std::endl;
		return nullptr;
	}
	
	return mesh_ptr;
}

std::shared_ptr<MeshData>
MeshData::LoadObjMeshData(const char* data_ptr, int data_size) {
	std::stringstream stream(data_ptr);
	auto mesh_ptr = std::make_shared<MeshData>();

	if (!mesh_ptr->LoadObjMesh(stream)) {
		std::cerr << "[ERROR] fail to parse mesh data" << std::endl;
		return nullptr;
	}

	return mesh_ptr;
}

std::shared_ptr<MeshData>
MeshData::LoadPlyMeshData(const char* data_ptr, int data_size) {
	std::stringstream stream(data_ptr);
	auto mesh_ptr = std::make_shared<MeshData>();

	if (!mesh_ptr->LoadPlyMesh(stream)) {
		std::cerr << "[ERROR] fail to parse mesh data" << std::endl;
		return nullptr;
	}

	return mesh_ptr;
}

bool MeshData::LoadObjMesh(std::istream& stm) {
	tinyobj::attrib_t attributes;
	std::vector<tinyobj::shape_t> shapes;
	std::vector<tinyobj::material_t> materials;
	std::string warning;
	std::string error;

	if (!tinyobj::LoadObj(
		&attributes, &shapes, &materials, &warning, &error, &stm, nullptr, true, false))
	{
		std::cerr << "TinyObjLoader failed to load data from buffer " << std::endl;
		return false;
	}

	size_t n_vertices = attributes.vertices.size() / 3;
	vertices.resize(n_vertices);
	const size_t n_vertices_bytes = attributes.vertices.size() * sizeof(float);
	std::memcpy(vertices.data(), attributes.vertices.data(), n_vertices_bytes);

	indices.clear();
	n_vertices = 0;
	for (const auto& shape : shapes)
		n_vertices += shape.mesh.num_face_vertices.size();
	indices.reserve(n_vertices);

	for (const auto& shape : shapes) {
		size_t index_offset = 0;
		for (size_t f = 0; f < shape.mesh.num_face_vertices.size(); ++f) {
			if (shape.mesh.num_face_vertices[f] != 3) {
				std::cerr << "Mesh contains non triangle shapes" << std::endl;
				index_offset += shape.mesh.num_face_vertices[f];
				continue;
			}

			bool geometry_counterclockwise_ = true;
			if (geometry_counterclockwise_) {
				indices.push_back(std::array<int, 3>{
					shape.mesh.indices[index_offset].vertex_index,
						shape.mesh.indices[index_offset + 1].vertex_index,
						shape.mesh.indices[index_offset + 2].vertex_index});
			}
			else {
				indices.push_back(std::array<int, 3>{
					shape.mesh.indices[index_offset + 2].vertex_index,
						shape.mesh.indices[index_offset + 1].vertex_index,
						shape.mesh.indices[index_offset].vertex_index});
			}
			index_offset += 3;
		}
	}
	return true;
}

bool MeshData::LoadPlyMesh(std::istream& stm) {
	tinyply::PlyFile file;
	file.parse_header(stm);

	std::shared_ptr<tinyply::PlyData> ply_verts;
	ply_verts = file.request_properties_from_element("vertex", { "x", "y", "z" });
	std::shared_ptr<tinyply::PlyData> ply_faces;
	ply_faces = file.request_properties_from_element("face", { "vertex_indices" }, 3);
	file.read(stm);

	if (ply_verts) {
		vertices.resize(ply_verts->count);
		const size_t n_vertices_bytes = ply_verts->buffer.size_bytes();
		std::memcpy(vertices.data(), ply_verts->buffer.get(), n_vertices_bytes);
	}

	if (ply_faces) {
		indices.resize(ply_faces->count);
		const size_t n_faces_bytes = ply_faces->buffer.size_bytes();
		std::memcpy(indices.data(), ply_faces->buffer.get(), n_faces_bytes);
	}

	return true;
}

std::shared_ptr<MeshData> MeshData::scaled(float scale) {
	auto scaled_mesh = std::make_shared<MeshData>();
	scaled_mesh->vertices.resize(vertices.size());
	for (int i = 0; i < vertices.size(); ++i) {
		scaled_mesh->vertices[i][0] = vertices[i][0] * scale;
		scaled_mesh->vertices[i][1] = vertices[i][1] * scale;
		scaled_mesh->vertices[i][2] = vertices[i][2] * scale;
	}
	scaled_mesh->indices = indices;
	return scaled_mesh;
}

BodyBuilder& BodyBuilder::obj_mesh_data(const char* data_ptr, int data_size, float scale) {
	std::stringstream stream(data_ptr);
	auto mesh_ptr = MeshData::LoadObjMeshData(data_ptr, data_size);
	if (nullptr == mesh_ptr) {
		delete body_ptr;
		body_ptr = nullptr;
	} else {
		if (scale != 1.0f) {
			mesh_ptr = mesh_ptr->scaled(scale);
		}
		body_ptr->mesh_ = mesh_ptr;
	}
	return *this;
}

BodyBuilder& BodyBuilder::ply_mesh_data(const char* data_ptr, int data_size, float scale) {
	std::stringstream stream(data_ptr);
	auto mesh_ptr = MeshData::LoadPlyMeshData(data_ptr, data_size);
	if (nullptr == mesh_ptr) {
		delete body_ptr;
		body_ptr = nullptr;
	} else {
		if (scale != 1.0f) {
			mesh_ptr = mesh_ptr->scaled(scale);
		}
		body_ptr->mesh_ = mesh_ptr;
	}
	return *this;
}

BodyBuilder& BodyBuilder::mesh_path(const std::filesystem::path& file_path, float scale) {
	auto mesh_ptr = MeshData::LoadMeshFile(file_path);
	if (nullptr == mesh_ptr) {
		delete body_ptr;
		body_ptr = nullptr;
	} else {
		if (scale != 1.0f) {
			mesh_ptr = mesh_ptr->scaled(scale);
		}
		body_ptr->mesh_ = mesh_ptr;
	}

	return *this;
}

BodyBuilder& BodyBuilder::body2world_pose(const Pose& pose) {
	body_ptr->body2world_pose_ = pose;
	return *this;
}

BodyBuilder& BodyBuilder::world2body_pose(const Pose& pose) {
	body_ptr->world2body_pose_ = pose;
	return *this;
}

float BodyBuilder::ComputeBodyDiameter() {
	auto vertices = body_ptr->mesh()->vertices;
	auto norm = [](const std::array<float, 3>& vertex) {
		return std::sqrt(
			vertex[0] * vertex[0] +
			vertex[1] * vertex[1] +
			vertex[2] * vertex[2]);
	};

	float max_radius = 0.0f;
	for (const auto& vertex : vertices) {
		max_radius = std::max(max_radius, norm(vertex));
	}
	return 2.0f * max_radius;
}

Body* BodyBuilder::Build() {
	if (nullptr == body_ptr) {
		return nullptr;
	}

	body_ptr->max_body_diameter_ = ComputeBodyDiameter();

	body_ptr->id_ = ++body_count;
	return body_ptr;
}

void Body::set_body2world_pose(const Pose& pose) {
	body2world_pose_ = pose;
	world2body_pose_ = body2world_pose_.inverse();
}

void Body::set_world2body_pose(const Pose& pose) {
	world2body_pose_ = pose;
	body2world_pose_ = world2body_pose_.inverse();
}

struct float2 { float x, y; };
struct float3 { float x, y, z; };
struct double3 { double x, y, z; };
struct uint3 { uint32_t x, y, z; };
struct uint4 { uint32_t x, y, z, w; };

struct geometry {
	std::vector<float3> vertices;
	std::vector<float3> normals;
	std::vector<float2> texcoords;
	std::vector<uint3> triangles;
};

void WritePly(const std::string& filename, const std::vector<Eigen::Vector3f>& points) {
	using namespace tinyply;

	auto n_pts = points.size();

	std::vector<float3> vec3;

	vec3.resize(n_pts);
	for (int i = 0; i < n_pts; ++i) {
		vec3[i].x = points[i].x();
		vec3[i].y = points[i].y();
		vec3[i].z = points[i].z();
	}

	std::filebuf fb_ascii;
	fb_ascii.open(filename, std::ios::out);
	std::ostream outstream_ascii(&fb_ascii);
	if (outstream_ascii.fail()) throw std::runtime_error("failed to open " + filename);

	tinyply::PlyFile ply_file;

	ply_file.add_properties_to_element("vertex", { "x", "y", "z" },
		Type::FLOAT32, vec3.size(), reinterpret_cast<uint8_t*>(vec3.data()), Type::INVALID, 0);
#if 0
	ply_file.add_properties_to_element("vertex", { "nx", "ny", "nz" },
		Type::FLOAT32, cube.normals.size(), reinterpret_cast<uint8_t*>(cube.normals.data()), Type::INVALID, 0);

	ply_file.add_properties_to_element("vertex", { "u", "v" },
		Type::FLOAT32, cube.texcoords.size(), reinterpret_cast<uint8_t*>(cube.texcoords.data()), Type::INVALID, 0);

	ply_file.add_properties_to_element("face", { "vertex_indices" },
		Type::UINT32, cube.triangles.size(), reinterpret_cast<uint8_t*>(cube.triangles.data()), Type::UINT8, 3);
#endif
	ply_file.get_comments().push_back("generated by tinyply 2.3");

	// Write an ASCII file
	ply_file.write(outstream_ascii, false);
}